package cn.firegod.demo.sharding.transaction;

import cn.firegod.demo.sharding.datasource.BasedJavaCodeConfig;
import org.apache.shardingsphere.transaction.core.TransactionType;
import org.apache.shardingsphere.transaction.core.TransactionTypeHolder;

import javax.sql.DataSource;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Random;

public class BasedJavaCodeTransaction extends BasedJavaCodeConfig {

    private int userIdStart = 1;
    private int insertCount = 3;
    private final DataSource dataSource;

    public BasedJavaCodeTransaction() throws SQLException {
        dataSource = super.shardingDataSource();
    }

    public void setInsertCount(int insertCount) {
        this.insertCount = insertCount;
    }

    public void setUserIdStart(int userIdStart) {
        this.userIdStart = userIdStart;
    }

    public void distributionTransactionWithXA() throws SQLException {
        TransactionTypeHolder.set(TransactionType.XA);
        doInsertOrders();
    }

    public void distributionTransactionWithSEATA() throws SQLException {
        TransactionTypeHolder.set(TransactionType.BASE);
        doInsertOrders();
    }

    private void doInsertOrders() throws SQLException {
        Connection connection = null;
        try {
            connection = dataSource.getConnection();
            connection.setAutoCommit(false);
            for (int i = 0; i < insertCount; i++) {
                PreparedStatement ps = connection.prepareStatement("insert into `order` (user_id,fee) values(?,?)");
                ps.setLong(1, userIdStart + i);
                ps.setBigDecimal(2, BigDecimal.valueOf(new Random().nextDouble()).setScale(2, RoundingMode.HALF_UP));
                ps.executeUpdate();
            }
            connection.commit();
        } catch (SQLException exception) {
            connection.rollback();
        }
    }
}
